package zdream.utils;

import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.util.ArrayList;
import java.util.List;
import java.util.Objects;

/**
 * 从 Photo Shop 等软件导出的 PNG 中带有很多无效信息, 使得本身 1KB 的文件大小变成了 12KB+
 * 这个程序将这类 PNG 中无效信息全清理掉, 保留全部颜色和图像的原始信息.
 *
 * @author Zdream
 * @date 2024-01-05
 * @since 0.0.1
 */
public class PngImageCleaner {
	public static void main(String[] args) throws IOException {
		if (args.length != 2) {
			System.err.println("Usage: <inPath> <outPath>");
			System.exit(1);
			return;
		}
		new PngImageCleaner().clean(args[0], args[1]);
	}

	private void clean(String inPath, String outPath) {
		File dirFile = new File(cleanPath(inPath));
		File outFile = new File(cleanPath(outPath));

		if (dirFile.isDirectory()) {
			File[] files = dirFile.listFiles(f -> !f.isDirectory() && f.getName().endsWith(".png"));
			if (files == null) {
				System.err.printf("扫描 %s 目录失败%n", dirFile);
				return;
			}

			for (File f : files) {
				try {
					convert(f, outputFile(f, dirFile, outFile));
				} catch (IOException e) {
					System.err.printf("转换 %s 文件失败: %s%n", f, e.getMessage());
				}
			}
		} else {
			try {
				convert(dirFile, outFile);
			} catch (IOException e) {
				System.err.printf("转换 %s 文件失败: %s%n", dirFile, e.getMessage());
			}
		}
	}

	private static String cleanPath(String path) {
		if (path.startsWith("\"")) {
			path = path.substring(1);
		}
		if (path.endsWith("\"")) {
			path = path.substring(0, path.length() - 1);
		}
		return path;
	}

	private static File outputFile(File f, File dirPath, File outPath) {
		String outFileName = f.getName();
		if (dirPath.equals(outPath)) {
			if (outFileName.endsWith(".png")) {
				outFileName = outFileName.substring(0, outFileName.length() - 4) + ".out.png";
			}
			return Paths.get(dirPath.getPath(), outFileName).toFile();
		}
		return Paths.get(outPath.getPath(), outFileName).toFile();
	}

	/**
	 * 清理 Adobe PS 保存文件中的大量无效信息
	 * 请首先将图片保存成 PNG 格式、索引颜色, 再执行本程序
	 */
	public static void convert(File f, File outFileName) throws IOException {
		System.out.printf("尝试转换文件 %s -> %s%n", f.getAbsolutePath(), outFileName.getAbsolutePath());

		final byte[] bs = Files.readAllBytes(f.toPath());
		int fileSize = (int) f.length();
		int index = 8;

		System.out.printf("以下是 %s 文件扫描的块信息:%n", f.getAbsolutePath());

		List<Chunk> chunks = new ArrayList<>();
		while (index < fileSize) {
			Chunk chunk = new Chunk(bs, index);
			System.out.println("\t" + chunk);
			chunks.add(chunk);
			index += 12 + chunk.size;

			if (chunk.isEndBlock()) {
				break;
			}
		}

		// 写
		FileOutputStream out = new FileOutputStream(outFileName);

		out.write(bs, 0, 8);
		for (Chunk chunk : chunks) {
			if (!chunk.keep) {
				continue;
			}
			out.write(chunk.content, chunk.indexFrom, chunk.indexTo - chunk.indexFrom);
		}

		out.flush();
		out.close();
	}

	/**
	 * 每个块由这些内容组成:
	 *
	 * <li>长度, 4 bytes: {@link Chunk#size}</li>
	 * <li>类型, 4 bytes: {@link Chunk#title}</li>
	 * <li>数据, 可变大小</li>
	 * <li>CRC 校验, 4 bytes</li>
	 */
	public static class Chunk {
		public static final String TITLE_END = "IEND";

		public final String title;
		public String desc;
		public final int size;
		public final byte[] content; // 文件的内容, 是全部的
		public final int indexFrom, indexTo;

		public boolean keep = true;

		public Chunk(byte[] content, int indexFrom) {
			this.content = Objects.requireNonNull(content);
			this.indexFrom = indexFrom;
			size = bytes2int(content, indexFrom);
			this.indexTo = indexFrom + size + 12;
			title = bytes2chars(content, indexFrom + 4);
			fillDescription();
		}

		public void fillDescription() {
			switch (title) {
			case "iTXt":
				this.desc = "文本注释";
				this.keep = false;
				break;
			case "tEXt":
				this.desc = "文本信息数据块";
				this.keep = false;
				break;
			case "zTXt":
				this.desc = "压缩文本数据块";
				this.keep = false;
				break;
			case "iCCP":
				this.desc = "嵌入式 ICC 配置文件块";
				this.keep = false;
				break;
			case "IHDR":
				this.desc = "头部信息块";
				break;
			case "pHYs":
				this.desc = "物理像素尺寸数据块";
				break;
			case "gAMA":
				this.desc = "伽马数据块";
				break;
			case "cHRM":
				this.desc = "基色和白色点数据块";
				break;
			case "sRGB":
				this.desc = "标准 RPG 色彩空间块";
				break;
			case "sBIT": // 在PLTE和IDAT之前
				this.desc = "样本有效位数据块";
				break;
			case "IDAT":
				this.desc = "图像数据块";
				break;
			case "PLTE":
				this.desc = "调色板数据块";
				break;
			case "tRNS":
				this.desc = "图像透明数据块";
				break;
			case TITLE_END:
				this.desc = "结束符";
				break;
			default:
				this.desc = "未知块";
				break;
			}
		}

		@Override
		public String toString() {
			return String.format("%s (%s): size=%d (%d~%d)%s",
					title, desc, size, indexFrom, indexTo, keep ? "" : " (被舍弃)");
		}

		public boolean isEndBlock() {
			return TITLE_END.equals(title);
		}
	}

	public static int bytes2int(byte[] bs, int index) {
		return (bs[index] & 0xFF) << 24 |
		(bs[index + 1] & 0xFF) << 16 |
		(bs[index + 2] & 0xFF) << 8 |
		(bs[index + 3] & 0xFF);
	}

	public static String bytes2chars(byte[] bs, int index) {
		return String.format("%s%s%s%s",
		(char) bs[index], (char) bs[index + 1], (char) bs[index + 2], (char) bs[index + 3]);
	}
}
